"""
This script simulates autonomous driving in a parking lot environment using the Prius model.
It integrates A* and RRT* path planning, local trajectory refinement, and PID/MPC control to navigate
a vehicle from a starting position to a goal while avoiding obstacles. 

Key Features:
- Generates a simulated parking lot environment with obstacles and boundaries.
- Implements global and local path planning using A* and RRT*.
- Implements MPC for local trajectory planning and control.
- Uses a PID controller for trajectory tracking and motion control (optional).
- Visualizes the planning and execution process in real time.

Note:
The simulation is repeated multiple times for different starting and goal positions
to calculate the success rate of the overall navigation system.
"""

import numpy as np
from urdfenvs.urdf_common.urdf_env import UrdfEnv
from bicycle_model_self import BicycleModelSelf
from urdfenvs.scene_examples.obstacles import *
from mpscenes.obstacles import *
from Create_Environment import Generate_Walls, Generate_Obstacles, Generate_Obstacle_Centers, create_grid_map, Generate_GoalPos, update_vehicle, initialize_plot, Generate_Dynamic_Obstacles, update_vehicle_dynobs
import matplotlib.pyplot as plt
from Simple_FreeSpace import compute_free_space
from RRT_Star import *
from A_Star_Planning import plan_path_through_goals
import time
from PID_Controller import pid_controller_incremental
import csv
import pandas as pd
from MPC_Controller import Run_MPC
from tqdm import tqdm

def run_prius(pos_idx, n_steps=2000, render=True, visualize=True, Control_Type = 'PID', Dynamic_Obstacles = False, print_log = True):
    """
    Run the simulation in the urdfenv environment.

    Parameters:
    pos_idx: float
        Index of the start and goal position.
    n_steps : int
        Maximum simulation steps.
    render : bool
        3D visualization.
    visualize : bool
        2D visualization.
    Control_Type : char
        Controller type (PID/MPC).
    Dynamic_Obstacles: bool
        Whether the environment contains dynamic obstacles
    print_log: bool
        print debug information

    Returns:
    history: list
        Simulation history information from urdfenvs.
    total_time: float
        Total simulation time in second.
    planning_success: bool
        Whether the vehicle reaches the goal position or not.
    """

    history = []
    total_time = 0
    # Set Robot
    robots = [
        BicycleModelSelf(
            urdf='prius.urdf',
            mode="vel",
            scaling=1,
            wheel_radius=0.31265,
            wheel_distance=2.86,
            spawn_offset=np.array([0.0, 0.0, 0.0]),
            actuated_wheels=['front_right_wheel_joint', 'front_left_wheel_joint', 'rear_right_wheel_joint', 'rear_left_wheel_joint'],
            steering_links=['front_right_steer_joint', 'front_left_steer_joint'],
            facing_direction='-x'
        )
    ]

    env: UrdfEnv = UrdfEnv(dt=0.01, robots=robots, render=render)
    
    # Define boundary and margin for obstacles
    length = 50
    width = 80
    margin = 2.5
    Walls, Wall_Dict = Generate_Walls(length, width, Type='Parking Lot', print_log = print_log)
    for Wall in Walls:
        env.add_obstacle(Wall)

    # Define obstacles
    Obs_centers = Generate_Obstacle_Centers(length, width, n_obstacles=16, Type='Parking Lot')
    Obstacles , Obs_Dict = Generate_Obstacles(Obs_centers, Type='Parking Lot')
    for Obstacle in Obstacles:
        env.add_obstacle(Obstacle)
        
    # Dynamic_Obs
    Dynamic_Obs_Trajs = []
    if Dynamic_Obstacles == True and Control_Type == 'MPC':
        Dynamic_Obstacles, Dynamic_Obs_Dicts = Generate_Dynamic_Obstacles(n_steps, n_obstacles=12, dt=0.01, Type="Parking Lot")
        for Dynamic_Obs in Dynamic_Obs_Dicts:
            Dynamic_Obs_Traj = []
            for x,y,z in Dynamic_Obs['geometry']['trajectory']['controlPoints']:
                Dynamic_Obs_Traj.append([x,y])
            Dynamic_Obs_Trajs.append(Dynamic_Obs_Traj)
        for Obstacle in Dynamic_Obstacles:
            env.add_obstacle(Obstacle)
    # Integrate boundary and obstacles
    obstacles_list = []
    for obs in Obs_Dict:
        obstacles_list.append(obs)
    for walls in Wall_Dict:
        obstacles_list.append(walls)
    global_map = {'l': length, 'w': width, 'm': margin, 'o': obstacles_list}

    # Define Grid Resolution
    grid_size = 0.1
    grid_map, min_x, max_x, min_y, max_y = create_grid_map(Wall_Dict, Obs_Dict, grid_size=grid_size)

    freespace_map = compute_free_space(
        grid_map, grid_size, 2.34+1.4, 2.42-1.4, 1.05, min_x, max_x, min_y, max_y, num_theta=16)
    if visualize:   
        fig, ax = initialize_plot(grid_map, min_x, max_x, min_y, max_y)
    else:
        ax = []
    # Define first action
    action = np.array([0.0, 0.0])  # Adjust action as necessary
    
    # Define start and goal in the world frame
    goals, goal_index = Generate_GoalPos(pos_idx)
    # pos0 = np.array([-20, -35, np.pi / 4, 0.0])
    if goal_index in {0, 2, 5, 6, 8}:
        pos0 = np.array([-20, -35, np.pi / 2, 0.0])
    else:
        pos0 = np.array([-20, -35, 0.0, 0.0])   
    goal = np.array([goals[-1, 0], goals[-1, 1], goals[-1, 2], 0.0])
    if print_log:
        print(f'Pos {pos0}, index {goal_index}')
    
    if visualize:  
        update_vehicle(ax, goal[0:3])
        update_vehicle(ax, pos0[0:3])
    
    # Reset simulation environment
    ob = env.reset(pos=pos0[0:3])
    if print_log:
        print(ob)
    # ---------Planning---------
    planning_success = False
    path_astar = plan_path_through_goals(pos0, goals, freespace_map, min_x, max_x, min_y, max_y, grid_size, ax, visualize=visualize)
    
    path, path_x, path_y, found_path_x, found_path_y, found_path_yaw, path_length, found_path_curvature, waypoints_dict = rrt_star(global_map, path_astar,
                                                                                                    pos0, goal, save_pkl = False, print_log = print_log)

    path_rear = []
    found_path_rear = []
    found_path_x_rear = []
    found_path_y_rear = []

    dis = 1.45
    for w in range(len(found_path_x)):
        found_path_rear.append((found_path_x[w] - dis * math.cos(found_path_yaw[w]), 
                               found_path_y[w] - dis * math.sin(found_path_yaw[w]), 
                               found_path_yaw[w], 
                               found_path_curvature[w]))
        found_path_x_rear.append(found_path_x[w] - dis * math.cos(found_path_yaw[w]))
        found_path_y_rear.append(found_path_y[w] - dis * math.sin(found_path_yaw[w]))

    for w in range(1, len(path)):
        path_rear.append((path[w][0] - dis * math.cos(path[w][2]), 
                          path[w][1] - dis * math.sin(path[w][2])))


    if path:
        if print_log:
            print("Found path:")
        if found_path_x:
            planning_success = True
        if visualize:
            draw_path_rrt(ax, path_rear, path_x, path_y, found_path_x_rear, found_path_y_rear)
    else:
        if print_log:
            print("No path found.")
        return history, total_time, planning_success
    
    path_lead = []
    for x,y in zip(found_path_x,found_path_y):
        path_lead.append([x,y])

    # ---------Planning---------
    if print_log:
        print(f"Initial observation : {ob}")
    
    # Reconfigure camera if needed
    env.reconfigure_camera(10.0, -30.0, -45.01, (0, -20, 0))
    
    pre_StrAng = 0
    pre_Speed = 0
    total_time = 0
    
    with tqdm(total=n_steps, desc="Overall Progress", unit="step") as progress_bar:
        for i in range(n_steps):
            planning_success = False
            ob, _, terminated, _, info = env.step(action)
            history.append(ob)
            pos = ob['robot_0']['joint_state']['position']
            pos = CoG_to_RearAxle(pos, dis=1.45)
            if print_log:
                print('Current Location: ', pos)
                
            speed = ob['robot_0']['joint_state']['forward_velocity'][0]
                
            if not path:
                print('Unsuccessful because no path found.')
                env.close()
                return history, total_time, planning_success
                break
            dyn_obs_center = []
            for dyn_obs_traj in Dynamic_Obs_Trajs:
                dyn_obs_center.append(dyn_obs_traj[i])
            # ------Control------
            
            if Control_Type == 'PID':
                new_speed, StrAng, _, _ = pid_controller_incremental(pos[:2], pos[2], pre_StrAng, speed, found_path_rear)
            elif Control_Type == 'MPC':
                new_speed, StrAng, planned_pathx, planned_pathy = Run_MPC(pos, goal[:3], path_lead, 5, i, Wall_Dict, Obs_Dict, Dynamic_Obs_Traj, [pre_Speed, pre_StrAng], print_log = print_log)
            
            if i % 20 == 0 and visualize and Control_Type == 'PID':
                update_vehicle(ax, pos)
            elif i % 20 == 0 and visualize and Control_Type == 'MPC':
                update_vehicle_dynobs(ax, pos, dyn_obs_center, planned_pathx, planned_pathy)
                
            action = np.array([new_speed, StrAng])
            pre_StrAng = StrAng 
            pre_Speed = new_speed
            total_time += 0.01
            progress_bar.update(1)
            if np.abs(pos[0] - goal[0]) < 0.2 and np.abs(pos[1] - goal[1]) < 0.2:
                planning_success = True
                print('Success')
                env.close()
                return history, total_time, planning_success
                break
            
            if terminated:
                print(info)
                print('Unsuccess because collision.')
                env.close()
                return history, total_time, planning_success
                break
            
            if i == n_steps-1:
                print(info)
                print('Unsuccess because reach max steps.')
                env.close()
                return history, total_time, planning_success
                break
    


def world_to_grid(world_pos, min_x, max_x, min_y, max_y, grid_size):
    """
    Convert world coordinates to grid coordinates.

    Parameters:
    world_pos -- Tuple (x, y) representing the position in world coordinates.
    min_x -- Minimum x-coordinate of the map.
    max_x -- Maximum x-coordinate of the map.
    min_y -- Minimum y-coordinate of the map.
    max_y -- Maximum y-coordinate of the map.
    grid_size -- Size of each grid cell.

    Returns:
    Tuple (grid_x, grid_y) representing the position in grid coordinates.
    """
    # Convert to grid coordinates
    grid_x = int((world_pos[0] - min_x) / grid_size)
    grid_y = int((world_pos[1] - min_y) / grid_size)

    # Check for out-of-bound values and clamp them within valid ranges
    grid_x = max(0, min(grid_x, int((max_x - min_x) / grid_size)))
    grid_y = max(0, min(grid_y, int((max_y - min_y) / grid_size)))

    return grid_x, grid_y


def grid_to_world(grid_pos, min_x, min_y, grid_size):
    """
    Convert grid coordinates to world coordinates.

    Parameters:
    grid_pos -- Tuple (grid_x, grid_y) representing the position in grid coordinates.
    min_x -- Minimum x-coordinate of the map.
    min_y -- Minimum y-coordinate of the map.
    grid_size -- Size of each grid cell.

    Returns:
    Numpy array [x, y] representing the position in world coordinates.
    """
    # Convert to world coordinates
    x = min_x + grid_pos[0] * grid_size
    y = min_y + grid_pos[1] * grid_size

    return np.array([x, y])



def CoG_to_RearAxle(CoG_pos, dis):
    """
    Convert the position of the center of gravity (CoG) to the position of the rear axle.

    Parameters:
    CoG_pos -- Numpy array or list [x_cog, y_cog, theta_cog], where:
        x_cog -- x-coordinate of the center of gravity.
        y_cog -- y-coordinate of the center of gravity.
        theta_cog -- orientation (angle) of the center of gravity in radians.
    dis -- Distance from the center of gravity to the rear axle.

    Returns:
    Numpy array [x_r, y_r, theta_cog], where:
        x_r -- x-coordinate of the rear axle.
        y_r -- y-coordinate of the rear axle.
        theta_cog -- orientation (angle), remains the same as the CoG orientation.
    """
    x_cog = CoG_pos[0]
    y_cog = CoG_pos[1]
    theta_cog = CoG_pos[2]
    
    # Calculate rear axle position based on the CoG position and distance
    x_r = x_cog - dis * np.cos(theta_cog)
    y_r = y_cog - dis * np.sin(theta_cog)
    
    return np.array([x_r, y_r, theta_cog])


import json
def save_progress(run_details, output_file="run_details.json"):
    # Save to local file
    with open(output_file, "w") as f:
        json.dump(run_details, f, indent=4)
    print(f"Progress saved to {output_file}")

if __name__ == "__main__":
    # Set goal indexes
    pos_indices = [0, 4, 6]
    run_details = {}

    for i in pos_indices:
        print('Run index: ', i)
        success_count = 0
        total_count = 0
        run_details[i] = []

        # Run every index 3 times
        for run in range(3):
            print('Run: ', run)
            total_count += 1
            _, _, planning_success = run_prius(
                pos_idx=i, render=False, visualize=True, Control_Type='MPC', Dynamic_Obstacles=True, print_log=False
            )
            print(f"Planning success for Run {run}: {planning_success}")

            # Save running information for debug
            run_info = {
                "Run index": i,
                "Run": run,
                "planning_success": planning_success
            }
            run_details[i].append(run_info)

            if planning_success:
                success_count += 1

        # Calculate success rate
        success_rate = success_count / total_count
        run_details[i].append({"Success rate": success_rate})
        print(f"Success rate for pos_idx {i}: {success_rate:.2f}")

        # Save result
        save_progress(run_details)

    print("Final results saved.")

    